'''This modules provides functions and classes for loading and handling
two-dimensional arrays (TDAs) defined in OpenDocument Spreadsheet files
(ODS).

The first row in a source spreadsheet must contain column names,
beginning with 'ID' (case-sensitive). The second row specifies the data
type for that column. The ID column must have type 'int'. The remaining
rows contain the data of the table and can have non-consecutive ID
values.

Note that each spreadsheet in an ODS file is treated as a separate TDA.
A spreadsheet name may be prefixed with an underscore to have it
ignored by the parser.

The following data types are allowed for columns: bool, int, long, str,
bytes, float, comment. The comment type is merely a placeholder and
instructs the parser to ignore that column.

The ODS parser currently adheres to the Python 3 string format, so a
column with type 'str' will produce a unicode string instance, and a
column of type 'bytes' will yield a normal string instance in
Python 2.x.
'''
from __future__ import print_function
from __future__ import unicode_literals

import os
import zipfile
import collections
import logging
from xml.parsers import expat
from os import path

__all__ = [
    'LoadODS',
    'LoadODSDir',
    'Merge',
    'CTDA',
    ]

GLogger = logging.getLogger('kharlia.tda')

def LoadODS(file, sort=True):
    '''Load an ODS file and return a dict mapping CTDA instances to
    names.
    '''
    with zipfile.ZipFile(file) as z, z.open('content.xml') as content:
        reader = CODSContentReader(content)

    if len(reader.Tables) == 0:
        return {}

    tdas = reader.BuildTDAs()
    if sort:
        for tda in tdas.itervalues():
            tda.Sort()
    return tdas

def LoadODSDir(name):
    '''Load a directory of ODS files.

    Returns: tdas:dict, conflicts:list
    '''
    GLogger.info('Loading ODS directory: %s', name)
    assert isinstance(name, unicode)
    conflicts = []
    tdas = {}

    fnames = os.listdir(name)
    for fname in fnames:
        if not fname.lower().endswith('.ods'):
            continue
        fpath = path.join(name, fname)
        GLogger.info('Loading ODS: %s', fpath)
        ftdas = LoadODS(fpath)
        for tname, tda in ftdas.iteritems():
            if tname in tdas:
                GLogger.warning('  Name conflict: %s', tname)
                conflicts.append(tname)
            else:
                GLogger.debug('  Loaded: %s', tname)
                tdas[tname] = tda

    return tdas, conflicts

def Merge(tdas, basetda, baseprefix='TDA_Base_', sort=True, trunc=True,
          recurse=True):
    '''Merge a map of TDAs. Returns number of merges done.

    Arguments:
    tdas -- Map of TDAs (name -> tda)
    basetda -- The base TDA containing Sheet and Merge columns
    baseprefix -- If specified, TDAs with this prefix will be merged
        into the base
    sort -- If True, sort all TDA rows
    trunc -- If True, delete TDAs from the map that are not listed in
        the basetda
    recurse -- If True, call this function recursively to perform
        another merge until no more merges can be done. This is
        necessary because the act of merging the base TDA adds more
        candidates for merging
    '''
    GLogger.debug('Merging TDA map with %s items...', len(tdas))
    mnames = {}
    merges = []

    if baseprefix:
        GLogger.debug('Including base in merge: %s with prefix %s',
                      basetda.Name, baseprefix)
        mnames[basetda.Name] = baseprefix

    GLogger.debug('TDAs:')
    for tname in tdas.iterkeys():
        GLogger.debug('  %s', tname)

    # Get the names of TDAs that will be created by merging
    GLogger.debug('Merge names:')
    for cname, csheet, cmerge in basetda.IterColumns('Name', 'Sheet', 'Merge'):
        if cmerge and csheet:
            mnames[cname] = csheet
            GLogger.debug('  %s with prefix %s', cname, csheet)

    # Determine which TDAs need to be merged and which need to be deleted
    # if they're not in the list
    GLogger.debug('Merge candidates:')
    for tname, tda in tdas.iteritems():
        if tda is basetda:
            continue
        for mname, mprefix in mnames.iteritems():
            if tname.startswith(mprefix):
                merges.append((mname, tname, tda))
                GLogger.debug('  %s into %s', tname, mname)
                continue

    # Perform merging
    GLogger.debug('Merging:')
    for mname, tname, tda in merges:
        mbase = tdas.get(mname, None)
        if mbase is None:
            tdas[mname] = tda
            tda.Name = mname
            GLogger.debug('  %s into %s as base', tname, mname)
        else:
            mbase.Update(tda)
            GLogger.debug('  %s into %s by update', tname, mname)
        del tdas[tname]


    # Repeat the merge, because the base TDA is also mergable, which
    if recurse and len(merges) != 0:
        GLogger.debug('Doing recursive merge...')
        while Merge(tdas, basetda, baseprefix, sort=False, trunc=False,
                    recurse=False) != 0:
            pass
        GLogger.debug('Finished recursive merge')

    # Now that all merging is likely completed, remove all TDAs not
    # listed in the base TDA
    if trunc:
        validnames = set(basetda.IterColumn('Name'))
        validnames.add(basetda.Name)
        names = set(tdas.iterkeys())
        names.difference_update(validnames)
        if len(names) > 0:
            GLogger.debug('Truncated TDAs:')
            for name in names:
                GLogger.debug('  %s', name)
        for name in names:
            del tdas[name]

    # Optional sort to ensure rows are in order after merging
    if sort:
        GLogger.debug('Sorting TDAs...')
        for tda in tdas.itervalues():
            tda.Sort()

    GLogger.debug('Finished TDA merge with %s merges', len(merges))
    return len(merges)

class CODSContentReader(object):
    '''Reads an ODS content.xml file.'''

    NAME_ROW = 0
    TYPE_ROW = 1
    DATA_ROW = 2

    TDA_TYPES = {
        'int': int,
        'str': unicode,
        'bytes': str,
        'bool': bool,
        'float': float,
        'long': long
        }

    TYPE_PARSERS = {
        bool: 'ParseBool'
        }

    COMMENT_TYPE = 'comment'
    NONE_VALUE = '****'

    def __init__(self, xmlfile=None):
        self.Reset()
        if xmlfile:
            self.Read(xmlfile)

    def Reset(self):
        self.Tables = {}
        self.CurName = None
        self.CurTable = None
        self.CurRow = None
        self.InCell = False
        self.InAnnot = False
        self.ColRepeat = 0

    def Read(self, xmlfile):
        p = expat.ParserCreate()
        p.StartElementHandler = self.HandleElementStart
        p.EndElementHandler = self.HandleElementEnd
        p.CharacterDataHandler = self.HandleCharData
        p.ParseFile(xmlfile)

    def BuildTDAs(self, nonevalue=NONE_VALUE):
        '''Build TDA objects from the strings parsed from the
        spreadsheets.
        '''
        tdas = {}
        for name, rows in self.Tables.iteritems():
            # Must have column name and type rows always
            if len(rows) < 2:
                raise ValueError('TDA {} does not have have minimum rows'
                                 .format(name))
            # Check for an ID column
            if rows[self.NAME_ROW][0] != 'ID' or rows[self.TYPE_ROW][0] != 'int':
                raise ValueError('TDA {} does not have ID column'
                                 .format(name))
            tda = CTDA(name)
            
            # ID Column should be added automatically        
            assert tda.NameMap.get('ID', None) == 0 and tda.Types[0] is int
            assert len(rows[self.NAME_ROW]) == len(rows[self.TYPE_ROW])

            # Add default parser for ID column
            parsers = [int]
            for cname, ctype in zip(rows[self.NAME_ROW][1:], rows[self.TYPE_ROW][1:]):
                # Comment type has no parser
                if ctype == self.COMMENT_TYPE:
                    parsers.append(None)
                    continue
                # Get the type object for the type name
                rtype = self.TDA_TYPES.get(ctype, None)
                if not rtype:
                    raise ValueError('invalid column type {}'.format(ctype))
                # Try getting a special parser for this type, otherwise use
                # the type itself as the parser.
                parsername = self.TYPE_PARSERS.get(rtype, None)
                parsers.append(getattr(self, parsername) if parsername else rtype)
                tda.AddColumn(cname, rtype)

            # Now do the actual value parsing and add rows to the TDA
            for row in rows[self.DATA_ROW:]:
                values = []
                for value, parser in zip(row, parsers):
                    # If parser is None, it's a comment column, ignore it
                    if parser is None:
                        continue
                    # None is specified by a specific string, usually '****'
                    if value == nonevalue:
                        values.append(None)
                    else:
                        values.append(parser(value))
                tda.AddRow(values, overwrite=True)
            tdas[name] = tda
        return tdas

    @staticmethod
    def ParseBool(value):
        '''Special parser for bool values in the TDA.'''
        value = value.lower()
        if value == '1' or value.startswith('t') or value.startswith('y'):
            return True
        if value == '0' or value.startswith('f') or value.startswith('n'):
            return False
        raise TypeError('invalid bool value {}'.format(value))

    def HandleElementStart(self, name, attrib):
        if name == 'table:table':
            tname = attrib['table:name']
            # We ignore tables prefixed with an underscore by setting
            # CurName to None
            if not tname.startswith('_'):
                self.CurName = attrib['table:name']
                assert self.CurName
                self.CurTable = []
            else:
                self.CurName = None
                self.CurTable = None

        elif not self.CurName:
            return

        elif name == 'table:table-row':
            self.CurRow = []

        elif name == 'table:table-cell':
            self.InCell = True
            if 'table:number-columns-repeated' in attrib:
                self.ColRepeat = int(attrib['table:number-columns-repeated'])

        # Annotations are within cells and also define char data, which needs
        # to be ignored in HandleCharData()
        elif name == 'office:annotation':
            self.InAnnot = True

    def HandleElementEnd(self, name):
        if not self.CurName:
            return

        if name == 'table:table':
            ctable = self.CurTable
            cname = self.CurName
            if len(ctable) > 0:
                self.Tables[cname] = ctable
            self.CurTable = None
            self.CurName = None
            
        elif name == 'table:table-row':
            if len(self.CurRow) > 0:
                self.CurTable.append(self.CurRow)
            self.CurRow = None

        elif name == 'table:table-cell':
            self.InCell = False
            self.ColRepeat = 0

        elif name == 'office:annotation':
            self.InAnnot = False

    def HandleCharData(self, data):
        if not self.CurName or not self.InCell or self.InAnnot:
            return

        if self.ColRepeat < 2:
            self.CurRow.append(str(data))
        else:
            ndata = str(data)
            crow = self.CurRow
            for i in xrange(self.ColRepeat):
                crow.append(ndata)

class CTDA(object):
    '''A two-dimensional array (TDA) of data.

    The rows are uniquely identified by the first integer ID column,
    and optionally uniquely identified by a second string ID column.
    '''
    ID_NAME = 'ID' # Name of ID column
    ID_TYPE = int # Type of ID column
    ID_INDEX = 0 # Index of the column
    STRID_TYPE = unicode # Type required for a string ID column
    STRID_INDEX = 1 # Index required for a string ID column

    def __init__(self, name=None):
        self.Name = name
        self.Clear()

    def Clear(self, addid=True):
        '''Clear rows and columns from the table.'''
        self.NameMap = collections.OrderedDict()
        self.Types = []
        self.IDMap = {}
        self.StrIDMap = {}
        self.Rows = []
        self.HasStrID = False
        if addid:
            self.AddColumn(self.ID_NAME, int)

    @property
    def ColumnCount(self):
        '''Number of columns in the table.'''
        return len(self.NameMap)

    @property
    def IsReady(self):
        '''Whether the table is ready for rows to be inserted.'''
        return (self.ColumnCount != 0 and self.ID_NAME in self.NameMap and
                self.Types[0] is self.ID_TYPE)

    def __getitem__(self, id):
        if isinstance(id, tuple) and len(id) == 2:
            return self.GetRowItem(id[0], id[1])
        else:
            return self.GetRow(id)

    def GetRow(self, id):
        '''Retrieve a row tuple by an integer or string ID.'''
        assert self.IsReady
        if isinstance(id, self.ID_TYPE):
            return self.IDMap[id]
        if isinstance(id, self.STRID_TYPE):
            if not self.HasStrID:
                raise ValueError('TDA does not have a string ID column')
            return self.StrIDMap[id]
        else:
            raise TypeError('invalid row ID type {}'.format(type(id).__name__))

    def GetRowAsDict(self, id):
        '''Retrieve a row as an OrderedDict.'''
        return collections.OrderedDict(zip(self.NameMap, self.GetRow(id)))

    def GetRowItem(self, id, colname):
        '''Retrieve an item from a row by the row's integer or string
        ID, and the name of the column.
        '''
        colindex = self.NameMap[colname]
        row = self.GetRow(id)
        return row[colindex]

    def HasColumn(self, name):
        '''Check if the column name exists in this table.'''
        return name in self.NameMap

    def GetColumnIndex(self, colname):
        '''Get the index of a column.'''
        return self.NameMap[colname]

    def IterColumn(self, name):
        '''Iterate through a single column.'''
        index = self.GetColumnIndex(name)
        for row in self.Rows:
            yield row[index]

    def IterColumns(self, *names):
        '''Iterate through multiple columns.'''
        iters = map(self.IterColumn, names)
        while True:
            yield map(next, iters)

    def AddColumn(self, name, ctype):
        '''Add a column to the table.

        This is only allowed if there are no rows in the table. The
        first column must be the ID column, and if the second column
        id a string ID column, then strings will be usable as row IDs.
        '''
        # ID column is always first
        if len(self.NameMap) == 0 and (name != self.ID_NAME or ctype != self.ID_TYPE):
            raise ValueError('first column must be {}:{}'
                             .format(self.ID_NAME, self.ID_TYPE.__name__))
        
        if name in self.NameMap:
            raise ValueError('column named "{}" already exists'.format(name))

        # No need to add columns after rows have been added to the table
        if len(self.Rows) != 0:
            raise ValueError('cannot add columns if rows exist')

        index = len(self.NameMap)
        self.NameMap[name] = index
        self.Types.append(ctype)
        assert self.Types[index] is ctype

        # Check if we can build a string ID map with the second column
        if index == self.STRID_INDEX and ctype == self.STRID_TYPE:
            self.HasStrID = True

    def AddRow(self, iterable, overwrite=True):
        '''Add a row to the table.

        Only allowed if IsReady is True. Iterable must yield number of
        values equal to ColumnCount. If overwrite is False, a
        ValueError will be raised if a row with the same ID already
        exists.
        '''
        assert self.IsReady

        # Rows are immutable
        row = tuple(iterable)
        if len(row) != self.ColumnCount:
            raise ValueError('iterable yielded {} values, need {}'
                             .format(len(row), self.ColumnCount))

        for value, ctype in zip(row, self.Types):
            if value is not None and not isinstance(value, ctype):
                raise ValueError('row value {} is not an instance of {}'
                                 .format(value, ctype))
        
        rowid = row[self.ID_INDEX]
        if rowid is None:
            raise ValueError('row ID cannot be None')
        exrow = self.IDMap.get(rowid, None)

        if not overwrite and exrow is not None:
            raise ValueError('row with ID {} already exists'.format(rowid))

        if exrow is not None:
            rowindex = self.Rows.index(exrow)
            self.Rows[rowindex] = row
        else:  
            rowindex = len(self.Rows)
            self.Rows.append(row)
        self.IDMap[rowid] = row

        if self.HasStrID:
            strid = row[self.STRID_INDEX]
            if strid is not None:
                assert isinstance(strid, self.STRID_TYPE)
                self.StrIDMap[strid] = row

        return rowindex

    def Update(self, other):
        '''Update this TDA with rows from other. Columns must match.'''
        assert self.IsReady

        if self.NameMap != other.NameMap or self.Types != other.Types:
            raise ValueError('Other TDA columns do not match.')

        assert self.HasStrID == other.HasStrID

        for row in other.Rows:
            self.AddRow(row, overwrite=True)

    def Sort(self):
        '''Sort rows by ID.'''
        assert self.IsReady
        idx = self.ID_INDEX
        self.Rows.sort(lambda x, y: cmp(x[idx], y[idx]))